import os.path as osp
import random

import numpy as np
from torch.utils.data import Dataset
import cv2
import pickle as pkl

random.seed(1)

# len of excluded_spec_ind is 111, [0, 1, 3, 6, 8, 9, 15, 20, 33, 35, 36, 43, 44, 50, 54, 55, 56, 57, 63, 67, 68, 69, 71, 73, 76, 80, 88, 92, 98, 99, 100, 102, 106, 107, 108, 110, 112, 114, 115, 117, 119, 125, 127, 129, 135, 144, 146, 147, 148, 149, 150, 152, 154, 155, 158, 159, 160, 163, 165, 169, 174, 176, 179, 185, 186, 192, 195, 197, 201, 202, 211, 215, 217, 230, 232, 235, 237, 240, 244, 247, 251, 252, 256, 258, 259, 262, 263, 267, 274, 278, 280, 283, 284, 291, 292, 296, 301, 302, 307, 309, 310, 313, 320, 326, 329, 331, 334, 338, 341, 346, 349]
# len of excluded_gene_ind is 8, [0, 1, 2, 3, 4, 5, 6, 7]
cifar10_excluded_spec_035labels = [0, 1, 3, 6, 8, 9, 15, 20, 33, 35, 36, 43, 44, 50, 54, 55, 56, 57, 63, 67, 68, 69, 71, 73, 76, 80, 88, 92, 98, 99, 100, 102, 106, 107, 108, 110, 112, 114, 115, 117, 119, 125, 127, 129, 135, 144, 146, 147, 148, 149, 150, 152, 154, 155, 158, 159, 160, 163, 165, 169, 174, 176, 179, 185, 186, 192, 195, 197, 201, 202, 211, 215, 217, 230, 232, 235, 237, 240, 244, 247, 251, 252, 256, 258, 259, 262, 263, 267, 274, 278, 280, 283, 284, 291, 292, 296, 301, 302, 307, 309, 310, 313, 320, 326, 329, 331, 334, 338, 341, 346, 349]
cifar10_excluded_gene_035labels = []

# len of excluded_spec_ind is 95, [8, 16, 18, 23, 29, 36, 38, 40, 43, 47, 51, 52, 53, 55, 56, 60, 66, 75, 77, 84, 85, 95, 96, 97, 102, 110, 118, 124, 130, 132, 133, 136, 137, 140, 142, 150, 151, 152, 156, 157, 163, 166, 167, 170, 173, 175, 176, 177, 178, 187, 188, 191, 198, 199, 204, 210, 214, 216, 217, 220, 223, 226, 228, 229, 236, 238, 241, 243, 249, 252, 257, 258, 264, 268, 270, 271, 275, 284, 286, 291, 299, 302, 305, 308, 313, 315, 316, 325, 328, 332, 335, 340, 343, 345, 348]
# len of excluded_gene_ind is 11, [1, 2, 6, 7, 9, 11, 13, 14, 15, 16, 19]
cifar100_excluded_spec_035labels = [8, 16, 18, 23, 29, 36, 38, 40, 43, 47, 51, 52, 53, 55, 56, 60, 66, 75, 77, 84, 85, 95, 96, 97, 102, 110, 118, 124, 130, 132, 133, 136, 137, 140, 142, 150, 151, 152, 156, 157, 163, 166, 167, 170, 173, 175, 176, 177, 178, 187, 188, 191, 198, 199, 204, 210, 214, 216, 217, 220, 223, 226, 228, 229, 236, 238, 241, 243, 249, 252, 257, 258, 264, 268, 270, 271, 275, 284, 286, 291, 299, 302, 305, 308, 313, 315, 316, 325, 328, 332, 335, 340, 343, 345, 348]
cifar100_excluded_gene_035labels = []

class TieredImageNet(Dataset):
    def __init__(self, transform=None, data_dir='../datasets/tiered-imagenet/', setname='train', exclude_cifar=True,
                 id_type='cifar10', num_included_gene_classes=-1, included_gene_classes=[]):
        label_pkl = osp.join(data_dir, setname + '_labels.pkl')
        data_pkl = osp.join(data_dir, setname + '_images_png.pkl')
        self.transform = transform
        if not exclude_cifar:
            self.excluded_spec_labels = []
            self.excluded_gene_labels = []
        else:
            if id_type == 'cifar10':
                self.excluded_spec_labels = cifar10_excluded_spec_035labels
                self.excluded_gene_labels = cifar10_excluded_gene_035labels
            elif id_type == 'cifar100':
                self.excluded_spec_labels = cifar100_excluded_spec_035labels
                self.excluded_gene_labels = cifar100_excluded_gene_035labels
            else:
                raise ValueError(f"Unsupported '{id_type}'")

        self.label_specific = []
        self.label_general = []
        self.data = []
        f_labels = open(label_pkl, 'rb')
        data_labels = pkl.load(f_labels, encoding='bytes')

        self.included_gene_classes = []
        if len(included_gene_classes) > 0:
            self.included_gene_classes = included_gene_classes
        else:
            if num_included_gene_classes > 0:
                temp_gene_classes = list(set(range(0, len(data_labels["label_general_str"]))).difference(set(self.excluded_gene_labels)))
                self.included_gene_classes = random.sample(temp_gene_classes, num_included_gene_classes)
            else:
                self.included_gene_classes = list(range(0, len(data_labels["label_general_str"])))
        print('Info, selected general classes are:', self.included_gene_classes)

        # self.include_indx = np.zeros_like(data_labels["label_specific"], dtype=bool)
        f_images = open(data_pkl, 'rb')
        data_images = pkl.load(f_images)
        for i, spec in enumerate(data_labels["label_specific"]):
            gene_label = data_labels["label_general"][i]
            if spec not in self.excluded_spec_labels and gene_label not in self.excluded_gene_labels and gene_label in self.included_gene_classes:
                self.label_specific.append(spec)
                self.label_general.append(gene_label)
                im = cv2.imdecode(data_images[i], 1)
                self.data.append(im)
                # if i > 523: break
                # print('i:', i)

        # self.data = [None for i in range(0, len(self.label_specific))]
        # ii = 0
        # for i, spec in enumerate(data_labels["label_specific"]):
        #     if spec not in self.excluded_spec_labels:
        #         self.data[ii] = cv2.imdecode(data_images[i], 1)
        #         ii = ii + 1

        self.label_specific = np.array(self.label_specific)
        self.label_general = np.array(self.label_general)
        self.data = np.array(self.data)

        self.label_specific_str = data_labels["label_specific_str"]
        self.label_general_str=data_labels["label_general_str"]

        print('Total number:{}, total selected samples:{}'.format(len(data_labels["label_specific"]), len(self.data)))
        # print('label_general_str:', self.label_general_str)
        # print('label_specific_str', self.label_specific_str)
        f_labels.close()
        f_images.close()

    def get_selected_classes(self):
        return self.included_gene_classes

    def __len__(self):
        return len(self.data)

    def __getitem__(self, i):
        x, y_spec, y_gen = self.data[i], self.label_specific[i], self.label_general[i]
        if self.transform is not None:
            x = self.transform(x)
        return x, y_spec, y_gen


def exlude_imagenet():
    import nltk
    nltk.download('wordnet')
    from nltk.corpus import wordnet
    def get_similarity(word1, word2):
        synset1 = wordnet.synsets(word1)
        synset2 = wordnet.synsets(word2)

        if synset1 and synset2:
            similarity = synset1[0].path_similarity(synset2[0])
            return similarity
        else:
            return 0

    # id_labels = ['airplane', 'automobile', 'bird', 'cat', 'deer', 'dog', 'frog', 'horse', 'ship', 'truck']
    id_labels = ['beaver', 'dolphin', 'otter', 'seal', 'whale', 'fish', 'flatfish', 'ray', 'shark', 'trout',
                 'orchid', 'poppy', 'rose', 'sunflower', 'tulip', 'bottle', 'bowl', 'can', 'cup', 'plate', 'apple',
                 'mushroom', 'orange', 'pear', 'sweet_pepper', 'clock', 'computer_keyboard', 'lamp', 'telephone',
                 'television', 'bed', 'chair', 'couch', 'table', 'wardrobe', 'bee', 'beetle', 'butterfly',
                 'caterpillar', 'cockroach', 'bear', 'leopard', 'lion', 'tiger', 'wolf', 'bridge', 'castle',
                 'house', 'road', 'skyscraper', 'cloud', 'forest', 'mountain', 'plain', 'sea', 'camel', 'cattle',
                 'chimpanzee', 'elephant', 'kangaroo', 'fox', 'porcupine', 'possum', 'raccoon', 'skunk', 'crab',
                 'lobster', 'snail', 'spider', 'worm', 'baby', 'boy', 'girl', 'man', 'woman', 'crocodile',
                 'dinosaur', 'lizard', 'snake', 'turtle', 'hamster', 'mouse', 'rabbit', 'shrew', 'squirrel',
                 'maple', 'oak', 'palm', 'pine', 'willow', 'bicycle', 'bus', 'motorcycle', 'pickup_truck', 'train',
                 'lawn_mower', 'rocket', 'streetcar', 'tank', 'tractor']
    excluded_spec_ind = set()
    excluded_gene_ind = set()
    label_pkl = osp.join('D:/PycharmProjects/datasets/tiered-imagenet/train_labels.pkl')
    for raw_id_label in id_labels:
        # id_label = wn.synset(raw_id_label + '.n.01')
        with open(label_pkl, 'rb') as f_labels:
            data_labels = pkl.load(f_labels, encoding='bytes')
            label_specific_str = data_labels["label_specific_str"]
            for i, spec_str in enumerate(label_specific_str):
                simi = 0
                if ', ' in spec_str:
                    temp_strs = spec_str.split(', ')
                    for temp_str in temp_strs:
                        temp_str = temp_str.replace(' ', '_')
                        temp_simi = get_similarity(raw_id_label, temp_str)
                        if temp_simi > simi:
                            simi = temp_simi
                else:
                    spec_str = spec_str.replace(' ', '_')
                    simi = get_similarity(raw_id_label, spec_str)
                # 0.18 for cifar10 but 0.2 for cifar100
                if simi > 0.2:
                    print(f"The similarity between '{raw_id_label}' and '{spec_str}' is {simi}")
                    excluded_spec_ind.add(i)

            label_general_str = data_labels["label_general_str"]
            for i, gene_str in label_general_str.items():
                simi = 0
                if ', ' in gene_str:
                    temp_strs = gene_str.split(', ')
                    for temp_str in temp_strs:
                        temp_str = temp_str.replace(' ', '_')
                        temp_simi = get_similarity(raw_id_label, temp_str)
                        if temp_simi is not None:
                            if temp_simi > simi:
                                simi = temp_simi
                else:
                    gene_str = gene_str.replace(' ', '_')
                    simi = get_similarity(raw_id_label, gene_str)
                if simi > 0.2:
                    print(f"The similarity between '{raw_id_label}' and '{gene_str}' is {simi}")
                    excluded_gene_ind.add(i)

    print('len of excluded_spec_ind is {}, {}'.format(len(list(excluded_spec_ind)), list(excluded_spec_ind)))
    print('len of excluded_gene_ind is {}, {}'.format(len(list(excluded_gene_ind)), list(excluded_gene_ind)))


# if __name__ == '__main__':
#     exlude_imagenet()
    # # import torchvision.transforms as T
    # # out_transform_train = T.Compose([
    # #     T.ToTensor(),
    # #     T.ToPILImage(),
    # #     T.
    # #     T.RandomCrop(32, padding=4),
    # #     T.RandomHorizontalFlip(),
    # #     T.ToTensor()
    # # ])
    # setname='train'
    # out_kwargs = {}
    # train_ood_loader = torch.utils.data.DataLoader(
    #     TieredImageNet(data_dir='D:/PycharmProjects/datasets/tiered-imagenet/', setname=setname, id_type='cifar10', num_included_gene_classes=4),
    #     batch_size=128, shuffle=False, **out_kwargs)
    # train_ood_iter = enumerate(train_ood_loader)
    # for i in range(0, 10000):
    #     ii, data_xy = next(train_ood_iter)
    #     ii
